package com.bjpowernode.controller;

import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.ai.ollama.api.OllamaOptions;
import org.springframework.http.MediaType;
import org.springframework.http.codec.ServerSentEvent;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.time.Duration;

@RequestMapping("/ai/ollama")
@Slf4j
@RestController
public class OllamaController {


    @Resource
    private OllamaChatModel chatModel;

    @RequestMapping(value = "/chat")
    public Object chat(@RequestParam(value = "msg") String msg) {
        String called = chatModel.call(msg);
        log.info(called);
        return called;
    }

    @RequestMapping(value = "/chat2")
    public Object chat2(@RequestParam(value = "msg") String msg) {
        ChatResponse chatResponse = chatModel.call(new Prompt(msg,
                OllamaOptions.builder()
                        .model("deepseek-r1:7b")    // 使用那个大模型
                        .temperature(0.4D)          // 温度, 温度值越高,准确率下降, 温度值越低,准确率提高
                        .build()));
        log.info(chatResponse.getResult().getOutput().getContent());
        return chatResponse.getResult().getOutput();
    }

    @RequestMapping(path = "/stream", produces = MediaType.TEXT_EVENT_STREAM_VALUE + ";charset=UTF-8")
    public Flux<ServerSentEvent<String>> stream(@RequestParam(value = "msg") String msg) {
        Flux<ChatResponse> flux = chatModel.stream(new Prompt(msg,
                OllamaOptions.builder()
                        .model("deepseek-r1:7b")    // 使用那个大模型
                        .temperature(0.4D)          // 温度, 温度值越高,准确率下降, 温度值越低,准确率提高
                        .build()));
        return flux.map(rsp -> {
            String content = rsp.getResult().getOutput().getContent();
            return ServerSentEvent.<String>builder(content).build();
        }).delayElements(Duration.ofMillis(100));
    }
}
